import numpy as np
import random
from enum import Enum

class CellType(Enum):
    """单元格类型枚举"""
    EMPTY = 0
    WUMPUS = 1
    PIT = 2
    GOLD = 3
    START = 4

class Action(Enum):
    """动作枚举"""
    UP = 0
    RIGHT = 1
    DOWN = 2
    LEFT = 3
    GRAB = 4
    SHOOT_UP = 5
    SHOOT_RIGHT = 6
    SHOOT_DOWN = 7
    SHOOT_LEFT = 8
    CLIMB = 9

class WumpusWorld:
    """Wumpus怪兽世界环境"""
    
    def __init__(self, size=4, pit_prob=0.2, random_positions=True):
        """
        初始化Wumpus世界
        
        参数:
        - size: 世界大小（默认4x4）
        - pit_prob: 每个单元格生成陷阱的概率
        - random_positions: 是否随机放置Wumpus和金子
        """
        self.size = size
        self.pit_prob = pit_prob
        self.random_positions = random_positions
        self.reset()
    
    def reset(self):
        """重置世界状态"""
        # 初始化空白世界
        self.grid = np.zeros((self.size, self.size), dtype=int)
        
        # 设置起点位置
        self.start_pos = (self.size-1, 0)
        self.grid[self.start_pos] = CellType.START.value
        
        # 放置Wumpus（不在起点）
        if self.random_positions:
            while True:
                wumpus_pos = (random.randint(0, self.size-1), random.randint(0, self.size-1))
                if wumpus_pos != self.start_pos:
                    break
        else:
            wumpus_pos = (0, 2)  # 固定位置
            
        self.wumpus_pos = wumpus_pos
        self.grid[wumpus_pos] = CellType.WUMPUS.value
        self.wumpus_alive = True
        
        # 放置金子（不在起点和Wumpus位置）
        if self.random_positions:
            while True:
                gold_pos = (random.randint(0, self.size-1), random.randint(0, self.size-1))
                if gold_pos != self.start_pos and gold_pos != wumpus_pos:
                    break
        else:
            gold_pos = (0, 3)  # 固定位置
            
        self.gold_pos = gold_pos
        self.grid[gold_pos] = CellType.GOLD.value
        
        # 放置陷阱（不在起点、Wumpus位置和金子位置）
        self.pit_positions = []
        for i in range(self.size):
            for j in range(self.size):
                pos = (i, j)
                if pos != self.start_pos and pos != wumpus_pos and pos != gold_pos:
                    if random.random() < self.pit_prob:
                        self.grid[pos] = CellType.PIT.value
                        self.pit_positions.append(pos)
        
        # 初始化代理状态
        self.agent_pos = self.start_pos
        self.has_gold = False
        self.has_arrow = True
        self.game_over = False
        self.won = False
        self.score = 0
        
        return self._get_state()
    
    def step(self, action):
        """
        执行一个动作并返回结果
        
        参数:
        - action: 要执行的动作
        
        返回:
        - next_state: 下一个状态
        - reward: 获得的奖励
        - done: 游戏是否结束
        - info: 额外信息
        """
        if self.game_over:
            return self._get_state(), 0, True, {"message": "游戏已经结束"}
        
        reward = -1  # 每一步都有-1的奖励（生存惩罚）
        info = {"message": ""}
        
        # 根据动作更新状态
        if action in [Action.UP.value, Action.RIGHT.value, Action.DOWN.value, Action.LEFT.value]:
            # 移动动作
            next_pos = self._get_next_position(action)
            if self._is_valid_position(next_pos):
                self.agent_pos = next_pos
                
                # 检查是否掉入陷阱或遇到Wumpus
                if self.grid[next_pos] == CellType.PIT.value:
                    reward = -1000  # 掉入陷阱惩罚
                    self.game_over = True
                    info["message"] = "代理掉入了陷阱!"
                elif self.grid[next_pos] == CellType.WUMPUS.value and self.wumpus_alive:
                    reward = -1000  # 被Wumpus杀死惩罚
                    self.game_over = True
                    info["message"] = "代理被Wumpus杀死!"
            else:
                reward = -10  # 尝试移动到无效位置惩罚
                info["message"] = "撞墙了!"
                
        elif action == Action.GRAB.value:
            # 拾取金子
            if self.agent_pos == self.gold_pos and not self.has_gold:
                self.has_gold = True
                reward = 1000  # 拾取金子奖励
                info["message"] = "代理获取了金子!"
            else:
                reward = -10  # 尝试拾取不存在的金子惩罚
                info["message"] = "这里没有金子!"
                
        elif action in [Action.SHOOT_UP.value, Action.SHOOT_RIGHT.value, 
                       Action.SHOOT_DOWN.value, Action.SHOOT_LEFT.value]:
            # 射箭
            if self.has_arrow:
                self.has_arrow = False
                shoot_dir = action - Action.SHOOT_UP.value  # 将射箭动作转换为方向
                if self._shoot_arrow(shoot_dir):
                    self.wumpus_alive = False
                    reward = 500  # 击杀Wumpus奖励
                    info["message"] = "代理击杀了Wumpus!"
                else:
                    reward = -10  # 射箭未命中惩罚
                    info["message"] = "箭未命中目标!"
            else:
                reward = -10  # 尝试使用已用完的箭惩罚
                info["message"] = "代理没有箭了!"
                
        elif action == Action.CLIMB.value:
            # 爬出洞穴
            if self.agent_pos == self.start_pos:
                if self.has_gold:
                    reward = 1000  # 成功完成任务奖励
                    self.game_over = True
                    self.won = True
                    info["message"] = "代理成功带着金子离开洞穴!"
                else:
                    reward = -100  # 未带金子离开惩罚
                    info["message"] = "代理未带金子就离开了洞穴!"
            else:
                reward = -10  # 尝试在非起点位置爬出惩罚
                info["message"] = "只能在起点位置爬出洞穴!"
        
        # 更新总分
        self.score += reward
        
        return self._get_state(), reward, self.game_over, info
    
    def _get_state(self):
        """获取当前状态表示"""
        # 状态表示为: (代理位置x, 代理位置y, 是否有金子, 是否有箭, Wumpus是否活着)
        return (self.agent_pos[0], self.agent_pos[1], 
                int(self.has_gold), int(self.has_arrow), int(self.wumpus_alive))
    
    def _get_next_position(self, action):
        """根据动作获取下一个位置"""
        x, y = self.agent_pos
        if action == Action.UP.value:
            return (x-1, y)
        elif action == Action.RIGHT.value:
            return (x, y+1)
        elif action == Action.DOWN.value:
            return (x+1, y)
        elif action == Action.LEFT.value:
            return (x, y-1)
        return self.agent_pos
    
    def _is_valid_position(self, pos):
        """检查位置是否有效"""
        x, y = pos
        return 0 <= x < self.size and 0 <= y < self.size
    
    def _shoot_arrow(self, direction):
        """射箭并检查是否击中Wumpus"""
        x, y = self.agent_pos
        wumpus_x, wumpus_y = self.wumpus_pos
        
        # 根据方向检查是否在同一行/列并且Wumpus在箭的路径上
        if direction == Action.UP.value - Action.SHOOT_UP.value:  # 向上射
            return y == wumpus_y and x > wumpus_x
        elif direction == Action.RIGHT.value - Action.SHOOT_UP.value:  # 向右射
            return x == wumpus_x and y < wumpus_y
        elif direction == Action.DOWN.value - Action.SHOOT_UP.value:  # 向下射
            return y == wumpus_y and x < wumpus_x
        elif direction == Action.LEFT.value - Action.SHOOT_UP.value:  # 向左射
            return x == wumpus_x and y > wumpus_y
        
        return False
    
    def get_percepts(self):
        """获取代理在当前位置的感知"""
        x, y = self.agent_pos
        
        # 检查是否有微风（陷阱附近）
        breeze = False
        for dx, dy in [(-1, 0), (0, 1), (1, 0), (0, -1)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.size and 0 <= ny < self.size:
                if self.grid[(nx, ny)] == CellType.PIT.value:
                    breeze = True
                    break
        
        # 检查是否有臭气（Wumpus附近）
        stench = False
        if self.wumpus_alive:
            wx, wy = self.wumpus_pos
            for dx, dy in [(-1, 0), (0, 1), (1, 0), (0, -1)]:
                nx, ny = wx + dx, wy + dy
                if 0 <= nx < self.size and 0 <= ny < self.size:
                    if (nx, ny) == (x, y):
                        stench = True
                        break
        
        # 检查是否有金光（金子位置）
        glitter = (x, y) == self.gold_pos and not self.has_gold
        
        return {
            "breeze": breeze,
            "stench": stench,
            "glitter": glitter,
            "bump": False,  # 碰撞在step函数中处理
            "scream": not self.wumpus_alive  # 如果Wumpus死了就会听到尖叫
        }